#include <iostream>
#include <string.h>
#include <mutex>
#include <unordered_map>
#include <stdio.h>
#include <unistd.h>
#include <vector>
#include <pthread.h>
#include <thread>
#include <queue>
#include "../buttonrpc/buttonrpc.hpp"
#include <ctime>
using namespace std;

#define TASK_TIMEOUT 10

struct Task{
    string fileName ;
    int IdMap;
// -------------
    int IdReduce;

    Task(string file, int id): fileName(file), IdMap(id){};
    Task(int id): IdReduce(id){};
};

// 超时检测的时间戳
struct TimeStamp{
    time_t time;  
    bool fin; // true 正常完成
};

// 回馈worker的消息
struct TaskResponse{
    int  state;             // 当前标识
    Task curTask;           // 当前任务
    int  curNumMapTask;     // 当前待分配map任务数量
    int  curNumReduceTask;  // 当前待分配reduce任务数量
    int  NumMapTask;        // 总map任务数量
    int  NumReduceTask;     // 总reduce任务数量
}; 


class Coordinator{

public:
    Coordinator(vector<string> files, int nReduce);
    TaskResponse AssignTask();                                    // 分配任务
    void*        TaskFin(int argsID);                             // 任务完成
    bool         TimeTick();                                      // 处理超时重发
    int          lenTaskFin(unordered_map<int, TimeStamp> &map);  // 获取正常完成任务的数量
    bool         Done();                                          // 任务全部完成
    ~Coordinator();

private:
    vector<string>    files;                       // 所有需要处理的文件——map的工作数量
    int               state;                       // 0 map 1 reduce 2 fin 标识位
    int               NumMapTask;                  // map任务数量
    int               NumReduceTask;               // reduce任务数量
    mutex             c_lock;                      // 互斥锁，保护共享数据
    queue<Task>       MapTasks;                    // 待处理的map工作队列
    queue<Task>       ReduceTasks;                 // 待处理的reduce工作队列
    unordered_map<int, TimeStamp> MapTaskFins;     // 处理完的map
    unordered_map<int, TimeStamp> ReduceTaskFins;  // 处理完的reduce

};

Coordinator::Coordinator(vector<string> files, int nReduce) : state(0), 
                                                              NumMapTask(files.size()),
                                                              NumReduceTask(nReduce),
                                                              files(files)
{
    time_t time_now = time(nullptr);
    // 分配map任务
    for (int i = 0; i < files.size(); i++)
    {
        this->MapTasks.emplace(Task(files[i], i));
        this->MapTaskFins.insert(make_pair(i, TimeStamp{time_now, false}));
    }
}

Coordinator::~Coordinator(){
    c_lock.~mutex();
}

int Coordinator::lenTaskFin(unordered_map<int, TimeStamp> &map){
    int len = 0;
    for (auto &&m : map){
        if (m.second.fin) len++;
    }
    return len;
}

TaskResponse Coordinator::AssignTask(){
    TaskResponse* reply;

    lock_guard<mutex> lock(c_lock); // 分配过程保证线程安全
    int staus = this->state;
    // map任务
    if (staus == 0){
        if (!this->MapTasks.empty()){
            Task maptask = MapTasks.front(); MapTasks.pop();
            reply->curTask = maptask;
            reply->curNumMapTask = MapTasks.size();
        } else{
            reply->curNumMapTask = -1;
        }
        reply->curNumReduceTask = ReduceTasks.size();
    } else if (staus == 1){ // reduce任务
        if (!this->ReduceTasks.empty()){
            Task reducetask = ReduceTasks.front(); ReduceTasks.pop();
            reply->curTask = reducetask;
            reply->curNumReduceTask = ReduceTasks.size();
        } else{
            reply->curNumReduceTask = -1;
        }
        reply->curNumMapTask = -1;
    }
    reply->state = staus;
    reply->NumReduceTask = this->NumReduceTask;
    reply->NumMapTask = this->NumMapTask;
    return *reply;
}

// argsID -> 任务id
void* Coordinator::TaskFin(int argsID){
    time_t time_now = time(nullptr);
    lock_guard<mutex> lock(c_lock);
    // 对 map 任务完成情况进行处理
    if (this->MapTaskFins.size() != this->NumMapTask){
        int start_time = MapTaskFins[argsID].time;
        if (time_now - start_time > TASK_TIMEOUT) return nullptr;   // 超时，丢弃数据
        this->MapTaskFins[argsID] = TimeStamp{time_now, true}; // 更新fin

        // map任务全部完成,开始分配reduce任务
        if (lenTaskFin(this->MapTaskFins) == this->NumMapTask){
            this->state = 1;
            for (int i = 0; i < this->NumReduceTask; i++)
            {
                this->ReduceTasks.emplace(Task(i));
                this->ReduceTaskFins[i] = TimeStamp{time_now, false};
            }
        }
    } else if (this->ReduceTaskFins.size() != this->NumReduceTask){     // 对 reduce 任务完成情况进行处理
        int start_time = ReduceTaskFins[argsID].time;
        if (time_now - start_time > TASK_TIMEOUT) return nullptr;       // 超时，丢弃数据
        this->ReduceTaskFins[argsID] = TimeStamp{time_now, true};  // 更新fin

        // 检查reduce任务是否全部完成
        if (lenTaskFin(this->MapTaskFins) == this->NumMapTask){
            this->state = 2;
        }
    }
    return nullptr;
}

/*
 *@brief 超时重发机制优化思考：改变现有Coordinator的分配工作队列和完成工作队列 $Tasks、$TaskFins
 *       更改方式：分配工作队列 $Tasks 不变；完成工作队列 $TaskFins 更换成正在处理的工作队列 $RunningTasks 和已完成工作数量 $FinNum，
 *       $RunningTasks队列会随着完成的工作减少size，在对 $RunningTasks 进行超时检测时的扫描时间复杂度会从O(n)降为O(n-m)——n为总任务数量,m为完成的任务数量。
 */
// 未优化版本
bool Coordinator::TimeTick(){
    lock_guard<mutex> lock(c_lock);
    
    time_t time_now = time(nullptr);
    int staus = this->state;

    if (staus == 0){
        // 检查是否有未完成的map任务已超时，若超时则重新加入工作队列
        for (int i = 0; i < this->NumMapTask; i++){
            TimeStamp tmp = this->MapTaskFins[i];
            if (!tmp.fin && time_now - tmp.time > TASK_TIMEOUT){
                cout << "map time out!" << endl;
                this->MapTasks.emplace(Task(this->files[i], i));     // 重新加入map工作队列
                this->MapTaskFins[i] = TimeStamp{time_now, false};   // 更新time
                return false;
            }
        }
    } else if (staus == 1){
        // 检查是否有未完成的reduce任务已超时，若超时则重新加入工作队列
        for (int i = 0; i < this->NumReduceTask; i++){
            TimeStamp tmp = this->ReduceTaskFins[i];
            if (!tmp.fin && time_now - tmp.time > TASK_TIMEOUT){
                cout << "map time out!" << endl;
                this->ReduceTasks.emplace(Task(this->files[i], i));     // 重新加入reduce工作队列
                this->ReduceTaskFins[i] = TimeStamp{time_now, false};   // 更新time
                return false;
            }
        }
    }
    return true;
}

bool Coordinator::Done(){
    if (!TimeTick())  return false;  // 若触发超时重发，则代表没有完成全部任务

    bool ret = false;
    if (lenTaskFin(this->ReduceTaskFins) == this->NumReduceTask){
        ret = true;
        lock_guard<mutex> lock(c_lock);
        this->state = 2;
    }
    return ret;
}

int main(){
    buttonrpc server;
    server.as_server(5555);

    vector<string> filenames = {
        "pg-being_ernest.txt",
        "pg-dorian_gray.txt",
        "pg-frankenstein.txt",
        "pg-grimm.txt",
        "pg-huckleberry_finn.txt",
        "pg-metamorphosis.txt",
        "pg-sherlock_holmes.txt",
        "pg-tom_sawyer.txt"
    };

    Coordinator master(filenames, 3);
    server.bind("AssignTask", &Coordinator::AssignTask, &master);
    server.bind("TaskFin", &Coordinator::TaskFin, &master);
    server.run();

    while (1)
    {
        if (master.Done()) break;
        sleep(3);
    }
    
    return 0;
}